package com.example.calcite.xml;

import org.apache.calcite.DataContext;
import org.apache.calcite.linq4j.Enumerable;
import org.apache.calcite.linq4j.Linq4j;
import org.apache.calcite.rel.type.RelDataType;
import org.apache.calcite.rel.type.RelDataTypeFactory;
import org.apache.calcite.schema.ScannableTable;
import org.apache.calcite.schema.impl.AbstractTable;
import org.w3c.dom.Document;
import org.w3c.dom.Element;
import org.w3c.dom.NodeList;

import java.util.ArrayList;
import java.util.List;

public class XmlTable extends AbstractTable implements ScannableTable {
    private final Document document;
    private final String collectionNodeName;
    private RelDataType rowType;

    public XmlTable(Document document, String collectionNodeName) {
        this.document = document;
        this.collectionNodeName = collectionNodeName;
    }

    @Override
    public RelDataType getRowType(RelDataTypeFactory typeFactory) {
        if (rowType == null) {
            NodeList elements = document.getElementsByTagName(collectionNodeName);
            if (elements.getLength() > 0) {
                Element firstElement = (Element) elements.item(0);
                List<RelDataType> types = new ArrayList<>();
                List<String> names = new ArrayList<>();
                
                // 添加属性
                if (firstElement.hasAttribute("id")) {
                    types.add(typeFactory.createJavaType(Integer.class));
                    names.add("id");
                }
                
                // 添加子元素
                NodeList children = firstElement.getChildNodes();
                for (int i = 0; i < children.getLength(); i++) {
                    if (children.item(i).getNodeType() == Element.ELEMENT_NODE) {
                        Element child = (Element) children.item(i);
                        types.add(typeFactory.createJavaType(String.class));
                        names.add(child.getNodeName());
                    }
                }
                
                rowType = typeFactory.createStructType(types, names);
            }
        }
        return rowType;
    }

    @Override
    public Enumerable<Object[]> scan(DataContext root) {
        NodeList elements = document.getElementsByTagName(collectionNodeName);
        List<Object[]> rows = new ArrayList<>();
        
        for (int i = 0; i < elements.getLength(); i++) {
            Element element = (Element) elements.item(i);
            List<Object> row = new ArrayList<>();
            
            // 添加属性
            if (element.hasAttribute("id")) {
                row.add(Integer.parseInt(element.getAttribute("id")));
            }
            
            // 添加子元素
            NodeList children = element.getChildNodes();
            for (int j = 0; j < children.getLength(); j++) {
                if (children.item(j).getNodeType() == Element.ELEMENT_NODE) {
                    Element child = (Element) children.item(j);
                    row.add(child.getTextContent());
                }
            }
            
            rows.add(row.toArray());
        }
        
        return Linq4j.asEnumerable(rows);
    }
}